# You need to set the environmental variable OPENAI_API_KEY
from langchain_core.tools import tool
from langchain import hub
from langchain_zhipu import ChatZhipuAI

from langchain.agents import create_react_agent
from langchain.agents import AgentType
from langchain.agents import AgentExecutor

from langchain_core.prompts import PromptTemplate

from logger import setup_logger
from langchain.agents.output_parsers import ReActJsonSingleInputOutputParser
from langchain.tools.render import render_text_description_and_args

from langchain_zhipu import ChatZhipuAI

from langchain_core.output_parsers import StrOutputParser
from langchain.prompts.chat import (ChatPromptTemplate, AIMessagePromptTemplate, SystemMessagePromptTemplate,
                                    HumanMessagePromptTemplate)

output_parser = StrOutputParser()

logger = setup_logger()
productInfoMap = {
    "1": {
        "id": "1",
        "name": "米家（MIJIA）米家台灯2Lite学生儿童学习专用阅读保护视力防蓝光读写台灯国AA",
        "url": "www.runoob.com",
        "price": 100
    },
    "2": {
        "id": "2",
        "name": "孩视宝立式护眼台灯大路灯LED全光谱落地学习灯儿童阅读灯钢琴灯E2-75W",
        "url": "www.jyshare.com",
        "price": 80
    },
    "3": {
        "id": "3",
        "name": "FSL佛山照明LED台灯学习护眼灯国AA级床头灯卧室宿舍阅读灯调色启明",
        "url": "www.google.com",
        "price": 30
    }
}


# Custom tool for the Agent

@tool
def product_recommendations_tool(product_name):
    """
    Recommend product methods based on product names.
    When this method returns the value as the final answer, please return the name, price, and purchase address all together。
    Input parameter: Product name
    Final Answer Output parameter: All product information
    Example:
    Input parameter: Table lamp
    Final Answer Output parameters:  {
            "productInfo": [
              {
                "id": "1",
                "name": "米家（MIJIA）米家台灯2Lite学生儿童学习专用阅读保护视力防蓝光读写台灯国AA",
                "url": "www.runoob.com",
                "price": 100
              },
              {
                "id": "2",
                "name": "孩视宝立式护眼台灯大路灯LED全光谱落地学习灯儿童阅读灯钢琴灯E2-75W",
                "url": "www.jyshare.com",
                "price": 80
              }
            ]
    }
    """
    print("product_recommendations_tool:" + product_name)

    return {"productInfo": [
        {"id": "1", "name": "米家（MIJIA）米家台灯2Lite学生儿童学习专用阅读保护视力防蓝光读写台灯国AA",
         "url": "www.runoob.com", "price": 100},
        {"id": "2", "name": "孩视宝立式护眼台灯大路灯LED全光谱落地学习灯儿童阅读灯钢琴灯E2-75W",
         "url": "www.jyshare.com", "price": 80},
        {"id": "3", "name": "FSL佛山照明LED台灯学习护眼灯国AA级床头灯卧室宿舍阅读灯调色启明", "url": "www.google.com",
         "price": 30}]}


llm = ChatZhipuAI(api_key="b9e2954e6e41c16b6c515203b35d1122.Vl0XqkeU4dNv0sjw"
                  , model="glm-4")  # 温度设置为0，结果随机性
template = """
你是客服AI助手，名字叫“小帮”。你的任务是针对用户的问题和要求提供专业、友好的服务。以下是你需要考虑的背景信息：[背景信息]。现在，用户提出了一个问题，请根据以下模板生成回复：

模板：
尊敬的用户，您好！我是[公司名称]的客服AI助手小帮。很抱歉给您带来不便。关于您的问题[用户问题]，我想了解以下细节：[需要补充的问题]。根据我的了解，[可能的解决方案]。请问您对此有何看法？如果还有其他问题，请随时告诉我。

注意事项：
0.不要假设或猜测，尽可能简短且友好的回复客户，等待客户提问
1. 保持语言礼貌、简洁。
2. 尽量提供明确的解决方案或建议。
3. 如果需要更多信息，请礼貌地询问。

现在，请根据以上模板回复以下用户问题：
"""
system_message_prompt = SystemMessagePromptTemplate.from_template(template)
human_template = """用户：“{text}”

补充信息（如果有的话）：
- 用户提到的产品/服务名称：[产品/服务名称]
- 用户遇到的问题类别：[问题类别]
- 用户的其他需求：[其他需求]"""
human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)

chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
# prompt = chat_prompt.from_messages(input_language="English", output_language="Chinese", text="How are you")
# @app.post("/chat")
# async def chat(request):
chain = chat_prompt | llm | output_parser


# Custom tool for the Agent
@tool
def get_answer(question):
    """
    Universal knowledge acquisition method, combining user questions with network knowledge base to answer questions
    """

    logger.info("get_answer:" + question)
    return chain.invoke(question)


# Saved React Prompt in langchain hub, we could manually type the prompt as well.
# prompt = hub.pull("hwchase17/react")

template = '''
SYSTEM

Answer the following questions as best you can. You have access to the following tools:



{tools}



The way you use the tools is by specifying a json blob.

Specifically, this json should have a `action` key (with the name of the tool to use) and a `action_input` key (with the input to the tool going here).



The only values that should be in the "action" field are: {tool_names}



The $JSON_BLOB should only contain a SINGLE action, do NOT return a list of multiple actions. Here is an example of a valid $JSON_BLOB:



```

{{

  "action": $TOOL_NAME,

  "action_input": $INPUT

}}

```



ALWAYS use the following format:



Question: the input question you must answer

Thought: you should always think about what to do

Action:

```

$JSON_BLOB

```

Observation: the result of the action

... (this Thought/Action/Observation can repeat N times)

Thought: I now know the final answer

Final Answer: the final answer to the original input question



Begin! Reminder to always use the exact characters `Final Answer` when responding.
If the tool returns the final answer, return it directly in the original information and its format, and do not omit the information returned by the tool.
Previous conversation history:

{chat_history}

New HUMAN input: {input}



{agent_scratchpad}'''

prompt = PromptTemplate.from_template(template)
model = ChatZhipuAI(api_key="b9e2954e6e41c16b6c515203b35d1122.Vl0XqkeU4dNv0sjw"
                    , model="glm-4", verbose=True)
model.do_sample = False  # 温度设置为0，结果随机性 ghbnm

tools = [product_recommendations_tool, get_answer]
agent = create_react_agent(
    model,
    tools,
    prompt,
    tools_renderer=render_text_description_and_args,
    output_parser=ReActJsonSingleInputOutputParser(),
)
# ,max_iterations=2
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True, handle_parsing_errors=True)
answer = agent_executor.invoke({"input": "我想买一个台灯有什么推荐？", "chat_history": None})
# answer = agent_executor.invoke({
#     "input": "都有什么区别?",
#     # Notice that chat_history is a string
#     # since this prompt is aimed at LLMs, not chat models
#     "chat_history": "Human: 我想买一个台灯有什么推荐？\nAI: 推荐的台灯产品如下：\n1. 米家（MIJIA）米家台灯2Lite，价格100元，购买地址www.runoob.com。\n2. 孩视宝立式护眼台灯E2-75W，价格80元，购买地址www.jyshare.com。\n3. FSL佛山照明LED台灯，价格30元，购买地址www.google.com。",
# })
# answer = agent_executor.invoke({"input": "台灯是W越大越好么？？","chat_history":None})
# answer = agent_executor.invoke({"input": "32厘米的锅有多大呢？？？","chat_history":None})
# answer = agent_executor.invoke({"input": "A4纸张有多大呢？？？","chat_history":"Human: 32厘米的锅有多大呢？？ \nAI: 32厘米的锅直径大约是32厘米，这意味着它的直径与一张标准A4纸的短边长度相近。这样的锅通常适合家庭烹饪，可以用来煮一些份量适中的食物。"})
print(answer)
